
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN"
  "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">

<html xmlns="http://www.w3.org/1999/xhtml" lang="en">
  <head>
    <meta http-equiv="X-UA-Compatible" content="IE=Edge" />
    <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
    <title>tensorflowonspark.TFNode &#8212; TensorFlowOnSpark 1.3.3 documentation</title>
    <link rel="stylesheet" href="../../_static/classic.css" type="text/css" />
    <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
    <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
    <script type="text/javascript" src="../../_static/jquery.js"></script>
    <script type="text/javascript" src="../../_static/underscore.js"></script>
    <script type="text/javascript" src="../../_static/doctools.js"></script>
    <link rel="index" title="Index" href="../../genindex.html" />
    <link rel="search" title="Search" href="../../search.html" /> 
  </head><body>
    <div class="related" role="navigation" aria-label="related navigation">
      <h3>Navigation</h3>
      <ul>
        <li class="right" style="margin-right: 10px">
          <a href="../../genindex.html" title="General Index"
             accesskey="I">index</a></li>
        <li class="right" >
          <a href="../../py-modindex.html" title="Python Module Index"
             >modules</a> |</li>
        <li class="nav-item nav-item-0"><a href="../../index.html">TensorFlowOnSpark 1.3.3 documentation</a> &#187;</li>
          <li class="nav-item nav-item-1"><a href="../index.html" accesskey="U">Module code</a> &#187;</li> 
      </ul>
    </div>  

    <div class="document">
      <div class="documentwrapper">
        <div class="bodywrapper">
          <div class="body" role="main">
            
  <h1>Source code for tensorflowonspark.TFNode</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2017 Yahoo Inc.</span>
<span class="c1"># Licensed under the terms of the Apache 2.0 license.</span>
<span class="c1"># Please see LICENSE file in the project root for terms.</span>
<span class="sd">&quot;&quot;&quot;This module provides helper functions for the TensorFlow application.</span>

<span class="sd">Primarily, these functions help with:</span>

<span class="sd">* starting the TensorFlow ``tf.train.Server`` for the node (allocating GPUs as desired, and determining the node&#39;s role in the cluster).</span>
<span class="sd">* managing input/output data for *InputMode.SPARK*.</span>
<span class="sd">&quot;&quot;&quot;</span>

<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">absolute_import</span>
<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">division</span>
<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">nested_scopes</span>
<span class="kn">from</span> <span class="nn">__future__</span> <span class="k">import</span> <span class="n">print_function</span>

<span class="kn">import</span> <span class="nn">getpass</span>
<span class="kn">import</span> <span class="nn">logging</span>
<span class="kn">import</span> <span class="nn">os</span>
<span class="kn">import</span> <span class="nn">time</span>
<span class="kn">from</span> <span class="nn">six.moves.queue</span> <span class="k">import</span> <span class="n">Empty</span>
<span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">marker</span>


<div class="viewcode-block" id="hdfs_path"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.hdfs_path">[docs]</a><span class="k">def</span> <span class="nf">hdfs_path</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">path</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Convenience function to create a Tensorflow-compatible absolute HDFS path from relative paths</span>

<span class="sd">  Args:</span>
<span class="sd">    :ctx: TFNodeContext containing the metadata specific to this node in the cluster.</span>
<span class="sd">    :path: path to convert</span>

<span class="sd">  Returns:</span>
<span class="sd">    An absolute path prefixed with the correct filesystem scheme.</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="c1">#  All Hadoop-Compatible File System Schemes (as of Hadoop 3.0.x):</span>
  <span class="n">HADOOP_SCHEMES</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;adl://&#39;</span><span class="p">,</span>
                    <span class="s1">&#39;hdfs://&#39;</span><span class="p">,</span>
                    <span class="s1">&#39;oss://&#39;</span><span class="p">,</span>
                    <span class="s1">&#39;s3://&#39;</span><span class="p">,</span>
                    <span class="s1">&#39;s3a://&#39;</span><span class="p">,</span>
                    <span class="s1">&#39;s3n://&#39;</span><span class="p">,</span>
                    <span class="s1">&#39;swift://&#39;</span><span class="p">,</span>
                    <span class="s1">&#39;viewfs://&#39;</span><span class="p">,</span>
                    <span class="s1">&#39;wasb://&#39;</span><span class="p">]</span>
  <span class="k">if</span> <span class="p">(</span><span class="nb">any</span><span class="p">(</span><span class="n">path</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="n">scheme</span><span class="p">)</span> <span class="k">for</span> <span class="n">scheme</span> <span class="ow">in</span> <span class="n">HADOOP_SCHEMES</span><span class="p">)</span>
      <span class="ow">or</span> <span class="n">path</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s1">&#39;file://&#39;</span><span class="p">)):</span>
    <span class="c1"># absolute path w/ scheme, just return as-is</span>
    <span class="k">return</span> <span class="n">path</span>
  <span class="k">elif</span> <span class="n">path</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;/&quot;</span><span class="p">):</span>
    <span class="c1"># absolute path w/o scheme, just prepend w/ defaultFS</span>
    <span class="k">return</span> <span class="n">ctx</span><span class="o">.</span><span class="n">defaultFS</span> <span class="o">+</span> <span class="n">path</span>
  <span class="k">else</span><span class="p">:</span>
    <span class="c1"># relative path, prepend defaultFS + standard working dir</span>
    <span class="k">if</span> <span class="n">ctx</span><span class="o">.</span><span class="n">defaultFS</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;hdfs://&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="n">ctx</span><span class="o">.</span><span class="n">defaultFS</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;viewfs://&quot;</span><span class="p">):</span>
      <span class="k">return</span> <span class="s2">&quot;</span><span class="si">{0}</span><span class="s2">/user/</span><span class="si">{1}</span><span class="s2">/</span><span class="si">{2}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ctx</span><span class="o">.</span><span class="n">defaultFS</span><span class="p">,</span> <span class="n">getpass</span><span class="o">.</span><span class="n">getuser</span><span class="p">(),</span> <span class="n">path</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">ctx</span><span class="o">.</span><span class="n">defaultFS</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;file://&quot;</span><span class="p">):</span>
      <span class="k">return</span> <span class="s2">&quot;</span><span class="si">{0}</span><span class="s2">/</span><span class="si">{1}</span><span class="s2">/</span><span class="si">{2}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ctx</span><span class="o">.</span><span class="n">defaultFS</span><span class="p">,</span> <span class="n">ctx</span><span class="o">.</span><span class="n">working_dir</span><span class="p">[</span><span class="mi">1</span><span class="p">:],</span> <span class="n">path</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
      <span class="n">logging</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;Unknown scheme </span><span class="si">{0}</span><span class="s2"> with relative path: </span><span class="si">{1}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ctx</span><span class="o">.</span><span class="n">defaultFS</span><span class="p">,</span> <span class="n">path</span><span class="p">))</span>
      <span class="k">return</span> <span class="s2">&quot;</span><span class="si">{0}</span><span class="s2">/</span><span class="si">{1}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ctx</span><span class="o">.</span><span class="n">defaultFS</span><span class="p">,</span> <span class="n">path</span><span class="p">)</span></div>


<div class="viewcode-block" id="start_cluster_server"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.start_cluster_server">[docs]</a><span class="k">def</span> <span class="nf">start_cluster_server</span><span class="p">(</span><span class="n">ctx</span><span class="p">,</span> <span class="n">num_gpus</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">rdma</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Function that wraps the creation of TensorFlow ``tf.train.Server`` for a node in a distributed TensorFlow cluster.</span>

<span class="sd">  This is intended to be invoked from within the TF ``map_fun``, replacing explicit code to instantiate ``tf.train.ClusterSpec``</span>
<span class="sd">  and ``tf.train.Server`` objects.</span>

<span class="sd">  Args:</span>
<span class="sd">    :ctx: TFNodeContext containing the metadata specific to this node in the cluster.</span>
<span class="sd">    :num_gpu: number of GPUs desired</span>
<span class="sd">    :rdma: boolean indicating if RDMA &#39;iverbs&#39; should be used for cluster communications.</span>

<span class="sd">  Returns:</span>
<span class="sd">    A tuple of (cluster_spec, server)</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="kn">import</span> <span class="nn">tensorflow</span> <span class="k">as</span> <span class="nn">tf</span>
  <span class="kn">from</span> <span class="nn">.</span> <span class="k">import</span> <span class="n">gpu_info</span>

  <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">{0}</span><span class="s2">: ======== </span><span class="si">{1}</span><span class="s2">:</span><span class="si">{2}</span><span class="s2"> ========&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ctx</span><span class="o">.</span><span class="n">worker_num</span><span class="p">,</span> <span class="n">ctx</span><span class="o">.</span><span class="n">job_name</span><span class="p">,</span> <span class="n">ctx</span><span class="o">.</span><span class="n">task_index</span><span class="p">))</span>
  <span class="n">cluster_spec</span> <span class="o">=</span> <span class="n">ctx</span><span class="o">.</span><span class="n">cluster_spec</span>
  <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">{0}</span><span class="s2">: Cluster spec: </span><span class="si">{1}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ctx</span><span class="o">.</span><span class="n">worker_num</span><span class="p">,</span> <span class="n">cluster_spec</span><span class="p">))</span>

  <span class="k">if</span> <span class="n">tf</span><span class="o">.</span><span class="n">test</span><span class="o">.</span><span class="n">is_built_with_cuda</span><span class="p">()</span> <span class="ow">and</span> <span class="n">num_gpus</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
    <span class="c1"># GPU</span>
    <span class="n">gpu_initialized</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="n">retries</span> <span class="o">=</span> <span class="mi">3</span>
    <span class="k">while</span> <span class="ow">not</span> <span class="n">gpu_initialized</span> <span class="ow">and</span> <span class="n">retries</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
      <span class="k">try</span><span class="p">:</span>
        <span class="c1"># override PS jobs to only reserve one GPU</span>
        <span class="k">if</span> <span class="n">ctx</span><span class="o">.</span><span class="n">job_name</span> <span class="o">==</span> <span class="s1">&#39;ps&#39;</span><span class="p">:</span>
          <span class="n">num_gpus</span> <span class="o">=</span> <span class="mi">1</span>

        <span class="c1"># Find a free gpu(s) to use</span>
        <span class="n">gpus_to_use</span> <span class="o">=</span> <span class="n">gpu_info</span><span class="o">.</span><span class="n">get_gpus</span><span class="p">(</span><span class="n">num_gpus</span><span class="p">)</span>
        <span class="n">gpu_prompt</span> <span class="o">=</span> <span class="s2">&quot;GPU&quot;</span> <span class="k">if</span> <span class="n">num_gpus</span> <span class="o">==</span> <span class="mi">1</span> <span class="k">else</span> <span class="s2">&quot;GPUs&quot;</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">{0}</span><span class="s2">: Using </span><span class="si">{1}</span><span class="s2">: </span><span class="si">{2}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ctx</span><span class="o">.</span><span class="n">worker_num</span><span class="p">,</span> <span class="n">gpu_prompt</span><span class="p">,</span> <span class="n">gpus_to_use</span><span class="p">))</span>

        <span class="c1"># Set GPU device to use for TensorFlow</span>
        <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;CUDA_VISIBLE_DEVICES&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">gpus_to_use</span>

        <span class="c1"># Create a cluster from the parameter server and worker hosts.</span>
        <span class="n">cluster</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">train</span><span class="o">.</span><span class="n">ClusterSpec</span><span class="p">(</span><span class="n">cluster_spec</span><span class="p">)</span>

        <span class="c1"># Create and start a server for the local task.</span>
        <span class="k">if</span> <span class="n">rdma</span><span class="p">:</span>
          <span class="n">server</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">train</span><span class="o">.</span><span class="n">Server</span><span class="p">(</span><span class="n">cluster</span><span class="p">,</span> <span class="n">ctx</span><span class="o">.</span><span class="n">job_name</span><span class="p">,</span> <span class="n">ctx</span><span class="o">.</span><span class="n">task_index</span><span class="p">,</span> <span class="n">protocol</span><span class="o">=</span><span class="s2">&quot;grpc+verbs&quot;</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
          <span class="n">server</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">train</span><span class="o">.</span><span class="n">Server</span><span class="p">(</span><span class="n">cluster</span><span class="p">,</span> <span class="n">ctx</span><span class="o">.</span><span class="n">job_name</span><span class="p">,</span> <span class="n">ctx</span><span class="o">.</span><span class="n">task_index</span><span class="p">)</span>
        <span class="n">gpu_initialized</span> <span class="o">=</span> <span class="kc">True</span>
      <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
        <span class="nb">print</span><span class="p">(</span><span class="n">e</span><span class="p">)</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">{0}</span><span class="s2">: Failed to allocate GPU, trying again...&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ctx</span><span class="o">.</span><span class="n">worker_num</span><span class="p">))</span>
        <span class="n">retries</span> <span class="o">-=</span> <span class="mi">1</span>
        <span class="n">time</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">10</span><span class="p">)</span>
    <span class="k">if</span> <span class="ow">not</span> <span class="n">gpu_initialized</span><span class="p">:</span>
      <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;Failed to allocate GPU&quot;</span><span class="p">)</span>
  <span class="k">else</span><span class="p">:</span>
    <span class="c1"># CPU</span>
    <span class="n">os</span><span class="o">.</span><span class="n">environ</span><span class="p">[</span><span class="s1">&#39;CUDA_VISIBLE_DEVICES&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="s1">&#39;&#39;</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">{0}</span><span class="s2">: Using CPU&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">ctx</span><span class="o">.</span><span class="n">worker_num</span><span class="p">))</span>

    <span class="c1"># Create a cluster from the parameter server and worker hosts.</span>
    <span class="n">cluster</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">train</span><span class="o">.</span><span class="n">ClusterSpec</span><span class="p">(</span><span class="n">cluster_spec</span><span class="p">)</span>

    <span class="c1"># Create and start a server for the local task.</span>
    <span class="n">server</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">train</span><span class="o">.</span><span class="n">Server</span><span class="p">(</span><span class="n">cluster</span><span class="p">,</span> <span class="n">ctx</span><span class="o">.</span><span class="n">job_name</span><span class="p">,</span> <span class="n">ctx</span><span class="o">.</span><span class="n">task_index</span><span class="p">)</span>

  <span class="k">return</span> <span class="p">(</span><span class="n">cluster</span><span class="p">,</span> <span class="n">server</span><span class="p">)</span></div>


<div class="viewcode-block" id="next_batch"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.next_batch">[docs]</a><span class="k">def</span> <span class="nf">next_batch</span><span class="p">(</span><span class="n">mgr</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">,</span> <span class="n">qname</span><span class="o">=</span><span class="s1">&#39;input&#39;</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;*DEPRECATED*. Use TFNode.DataFeed class instead.&quot;&quot;&quot;</span>
  <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;DEPRECATED: Use TFNode.DataFeed class instead&quot;</span><span class="p">)</span></div>


<div class="viewcode-block" id="export_saved_model"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.export_saved_model">[docs]</a><span class="k">def</span> <span class="nf">export_saved_model</span><span class="p">(</span><span class="n">sess</span><span class="p">,</span> <span class="n">export_dir</span><span class="p">,</span> <span class="n">tag_set</span><span class="p">,</span> <span class="n">signatures</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;Convenience function to export a saved_model using provided arguments</span>

<span class="sd">  The caller specifies the saved_model signatures in a simplified python dictionary form, as follows::</span>

<span class="sd">    signatures = {</span>
<span class="sd">      &#39;signature_def_key&#39;: {</span>
<span class="sd">        &#39;inputs&#39;: { &#39;input_tensor_alias&#39;: input_tensor_name },</span>
<span class="sd">        &#39;outputs&#39;: { &#39;output_tensor_alias&#39;: output_tensor_name },</span>
<span class="sd">        &#39;method_name&#39;: &#39;method&#39;</span>
<span class="sd">      }</span>
<span class="sd">    }</span>

<span class="sd">  And this function will generate the `signature_def_map` and export the saved_model.</span>

<span class="sd">  Args:</span>
<span class="sd">    :sess: a tf.Session instance</span>
<span class="sd">    :export_dir: path to save exported saved_model</span>
<span class="sd">    :tag_set: string tag_set to identify the exported graph</span>
<span class="sd">    :signatures: simplified dictionary representation of a TensorFlow signature_def_map</span>

<span class="sd">  Returns:</span>
<span class="sd">    A saved_model exported to disk at ``export_dir``.</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="kn">import</span> <span class="nn">tensorflow</span> <span class="k">as</span> <span class="nn">tf</span>
  <span class="n">g</span> <span class="o">=</span> <span class="n">sess</span><span class="o">.</span><span class="n">graph</span>
  <span class="n">g</span><span class="o">.</span><span class="n">_unsafe_unfinalize</span><span class="p">()</span>           <span class="c1"># https://github.com/tensorflow/serving/issues/363</span>
  <span class="n">builder</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">saved_model</span><span class="o">.</span><span class="n">builder</span><span class="o">.</span><span class="n">SavedModelBuilder</span><span class="p">(</span><span class="n">export_dir</span><span class="p">)</span>

  <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== signatures: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">signatures</span><span class="p">))</span>
  <span class="n">signature_def_map</span> <span class="o">=</span> <span class="p">{}</span>
  <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">sig</span> <span class="ow">in</span> <span class="n">signatures</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
    <span class="n">signature_def_map</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">tf</span><span class="o">.</span><span class="n">saved_model</span><span class="o">.</span><span class="n">signature_def_utils</span><span class="o">.</span><span class="n">build_signature_def</span><span class="p">(</span>
        <span class="n">inputs</span><span class="o">=</span><span class="p">{</span><span class="n">name</span><span class="p">:</span> <span class="n">tf</span><span class="o">.</span><span class="n">saved_model</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">build_tensor_info</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">tensor</span> <span class="ow">in</span> <span class="n">sig</span><span class="p">[</span><span class="s1">&#39;inputs&#39;</span><span class="p">]</span><span class="o">.</span><span class="n">items</span><span class="p">()},</span>
        <span class="n">outputs</span><span class="o">=</span><span class="p">{</span><span class="n">name</span><span class="p">:</span> <span class="n">tf</span><span class="o">.</span><span class="n">saved_model</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">build_tensor_info</span><span class="p">(</span><span class="n">tensor</span><span class="p">)</span> <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="n">tensor</span> <span class="ow">in</span> <span class="n">sig</span><span class="p">[</span><span class="s1">&#39;outputs&#39;</span><span class="p">]</span><span class="o">.</span><span class="n">items</span><span class="p">()},</span>
        <span class="n">method_name</span><span class="o">=</span><span class="n">sig</span><span class="p">[</span><span class="s1">&#39;method_name&#39;</span><span class="p">]</span> <span class="k">if</span> <span class="s1">&#39;method_name&#39;</span> <span class="ow">in</span> <span class="n">sig</span> <span class="k">else</span> <span class="n">key</span><span class="p">)</span>
  <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;===== signature_def_map: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">signature_def_map</span><span class="p">))</span>
  <span class="n">builder</span><span class="o">.</span><span class="n">add_meta_graph_and_variables</span><span class="p">(</span>
      <span class="n">sess</span><span class="p">,</span>
      <span class="n">tag_set</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;,&#39;</span><span class="p">),</span>
      <span class="n">signature_def_map</span><span class="o">=</span><span class="n">signature_def_map</span><span class="p">,</span>
      <span class="n">clear_devices</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
  <span class="n">g</span><span class="o">.</span><span class="n">finalize</span><span class="p">()</span>
  <span class="n">builder</span><span class="o">.</span><span class="n">save</span><span class="p">()</span></div>


<div class="viewcode-block" id="batch_results"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.batch_results">[docs]</a><span class="k">def</span> <span class="nf">batch_results</span><span class="p">(</span><span class="n">mgr</span><span class="p">,</span> <span class="n">results</span><span class="p">,</span> <span class="n">qname</span><span class="o">=</span><span class="s1">&#39;output&#39;</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;*DEPRECATED*. Use TFNode.DataFeed class instead.&quot;&quot;&quot;</span>
  <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;DEPRECATED: Use TFNode.DataFeed class instead&quot;</span><span class="p">)</span></div>


<div class="viewcode-block" id="terminate"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.terminate">[docs]</a><span class="k">def</span> <span class="nf">terminate</span><span class="p">(</span><span class="n">mgr</span><span class="p">,</span> <span class="n">qname</span><span class="o">=</span><span class="s1">&#39;input&#39;</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;*DEPRECATED*. Use TFNode.DataFeed class instead.&quot;&quot;&quot;</span>
  <span class="k">raise</span> <span class="ne">Exception</span><span class="p">(</span><span class="s2">&quot;DEPRECATED: Use TFNode.DataFeed class instead&quot;</span><span class="p">)</span></div>


<div class="viewcode-block" id="DataFeed"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.DataFeed">[docs]</a><span class="k">class</span> <span class="nc">DataFeed</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
  <span class="sd">&quot;&quot;&quot;This class manages the *InputMode.SPARK* data feeding process from the perspective of the TensorFlow application.</span>

<span class="sd">  Args:</span>
<span class="sd">    :mgr: TFManager instance associated with this Python worker process.</span>
<span class="sd">    :train_mode: boolean indicating if the data feed is expecting an output response (e.g. inferencing).</span>
<span class="sd">    :qname_in: *INTERNAL_USE*</span>
<span class="sd">    :qname_out: *INTERNAL_USE*</span>
<span class="sd">    :input_mapping: *For Spark ML Pipelines only*.  Dictionary of input DataFrame columns to input TensorFlow tensors.</span>
<span class="sd">  &quot;&quot;&quot;</span>
  <span class="k">def</span> <span class="nf">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">mgr</span><span class="p">,</span> <span class="n">train_mode</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">qname_in</span><span class="o">=</span><span class="s1">&#39;input&#39;</span><span class="p">,</span> <span class="n">qname_out</span><span class="o">=</span><span class="s1">&#39;output&#39;</span><span class="p">,</span> <span class="n">input_mapping</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>

    <span class="bp">self</span><span class="o">.</span><span class="n">mgr</span> <span class="o">=</span> <span class="n">mgr</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">train_mode</span> <span class="o">=</span> <span class="n">train_mode</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">qname_in</span> <span class="o">=</span> <span class="n">qname_in</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">qname_out</span> <span class="o">=</span> <span class="n">qname_out</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">done_feeding</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">input_tensors</span> <span class="o">=</span> <span class="p">[</span><span class="n">tensor</span> <span class="k">for</span> <span class="n">col</span><span class="p">,</span> <span class="n">tensor</span> <span class="ow">in</span> <span class="nb">sorted</span><span class="p">(</span><span class="n">input_mapping</span><span class="o">.</span><span class="n">items</span><span class="p">())]</span> <span class="k">if</span> <span class="n">input_mapping</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="kc">None</span>

<div class="viewcode-block" id="DataFeed.next_batch"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.DataFeed.next_batch">[docs]</a>  <span class="k">def</span> <span class="nf">next_batch</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Gets a batch of items from the input RDD.</span>

<span class="sd">    If multiple tensors are provided per row in the input RDD, e.g. tuple of (tensor1, tensor2, ..., tensorN) and:</span>

<span class="sd">    * no ``input_mapping`` was provided to the DataFeed constructor, this will return an array of ``batch_size`` tuples,</span>
<span class="sd">      and the caller is responsible for separating the tensors.</span>
<span class="sd">    * an ``input_mapping`` was provided to the DataFeed constructor, this will return a dictionary of N tensors,</span>
<span class="sd">      with tensor names as keys and arrays of length ``batch_size`` as values.</span>

<span class="sd">    Note: if the end of the data is reached, this may return with fewer than ``batch_size`` items.</span>

<span class="sd">    Args:</span>
<span class="sd">      :batch_size: number of items to retrieve.</span>

<span class="sd">    Returns:</span>
<span class="sd">      A batch of items or a dictionary of tensors.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;next_batch() invoked&quot;</span><span class="p">)</span>
    <span class="n">queue</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">qname_in</span><span class="p">)</span>
    <span class="n">tensors</span> <span class="o">=</span> <span class="p">[]</span> <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_tensors</span> <span class="ow">is</span> <span class="kc">None</span> <span class="k">else</span> <span class="p">{</span><span class="n">tensor</span><span class="p">:</span> <span class="p">[]</span> <span class="k">for</span> <span class="n">tensor</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_tensors</span><span class="p">}</span>
    <span class="n">count</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="k">while</span> <span class="n">count</span> <span class="o">&lt;</span> <span class="n">batch_size</span><span class="p">:</span>
      <span class="n">item</span> <span class="o">=</span> <span class="n">queue</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">block</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
      <span class="k">if</span> <span class="n">item</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># End of Feed</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;next_batch() got None&quot;</span><span class="p">)</span>
        <span class="n">queue</span><span class="o">.</span><span class="n">task_done</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">done_feeding</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="k">break</span>
      <span class="k">elif</span> <span class="nb">type</span><span class="p">(</span><span class="n">item</span><span class="p">)</span> <span class="ow">is</span> <span class="n">marker</span><span class="o">.</span><span class="n">EndPartition</span><span class="p">:</span>
        <span class="c1"># End of Partition</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;next_batch() got EndPartition&quot;</span><span class="p">)</span>
        <span class="n">queue</span><span class="o">.</span><span class="n">task_done</span><span class="p">()</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">train_mode</span> <span class="ow">and</span> <span class="n">count</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
          <span class="k">break</span>
      <span class="k">else</span><span class="p">:</span>
        <span class="c1"># Normal item</span>
        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_tensors</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
          <span class="n">tensors</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">item</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
          <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_tensors</span><span class="p">)):</span>
            <span class="n">tensors</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">input_tensors</span><span class="p">[</span><span class="n">i</span><span class="p">]]</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">item</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="n">count</span> <span class="o">+=</span> <span class="mi">1</span>
        <span class="n">queue</span><span class="o">.</span><span class="n">task_done</span><span class="p">()</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;next_batch() returning </span><span class="si">{0}</span><span class="s2"> items&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">count</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">tensors</span></div>

<div class="viewcode-block" id="DataFeed.should_stop"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.DataFeed.should_stop">[docs]</a>  <span class="k">def</span> <span class="nf">should_stop</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Check if the feed process was told to stop (by a call to ``terminate``).&quot;&quot;&quot;</span>
    <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">done_feeding</span></div>

<div class="viewcode-block" id="DataFeed.batch_results"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.DataFeed.batch_results">[docs]</a>  <span class="k">def</span> <span class="nf">batch_results</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">results</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Push a batch of output results to the Spark output RDD of ``TFCluster.inference()``.</span>

<span class="sd">    Note: this currently expects a one-to-one mapping of input to output data, so the length of the ``results`` array should match the length of</span>
<span class="sd">    the previously retrieved batch of input data.</span>

<span class="sd">    Args:</span>
<span class="sd">      :results: array of output data for the equivalent batch of input data.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;batch_results() invoked&quot;</span><span class="p">)</span>
    <span class="n">queue</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">qname_out</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">results</span><span class="p">:</span>
      <span class="n">queue</span><span class="o">.</span><span class="n">put</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="n">block</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">debug</span><span class="p">(</span><span class="s2">&quot;batch_results() returning data&quot;</span><span class="p">)</span></div>

<div class="viewcode-block" id="DataFeed.terminate"><a class="viewcode-back" href="../../tensorflowonspark.TFNode.html#tensorflowonspark.TFNode.DataFeed.terminate">[docs]</a>  <span class="k">def</span> <span class="nf">terminate</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Terminate data feeding early.</span>

<span class="sd">    Since TensorFlow applications can often terminate on conditions unrelated to the training data (e.g. steps, accuracy, etc),</span>
<span class="sd">    this method signals the data feeding process to ignore any further incoming data.  Note that Spark itself does not have a mechanism</span>
<span class="sd">    to terminate an RDD operation early, so the extra partitions will still be sent to the executors (but will be ignored).  Because</span>
<span class="sd">    of this, you should size your input data accordingly to avoid excessive overhead.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;terminate() invoked&quot;</span><span class="p">)</span>
    <span class="bp">self</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="s1">&#39;state&#39;</span><span class="p">,</span> <span class="s1">&#39;terminating&#39;</span><span class="p">)</span>

    <span class="c1"># drop remaining items in the queue</span>
    <span class="n">queue</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">mgr</span><span class="o">.</span><span class="n">get_queue</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">qname_in</span><span class="p">)</span>
    <span class="n">count</span> <span class="o">=</span> <span class="mi">0</span>
    <span class="n">done</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="k">while</span> <span class="ow">not</span> <span class="n">done</span><span class="p">:</span>
      <span class="k">try</span><span class="p">:</span>
        <span class="n">queue</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">block</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">timeout</span><span class="o">=</span><span class="mi">5</span><span class="p">)</span>
        <span class="n">queue</span><span class="o">.</span><span class="n">task_done</span><span class="p">()</span>
        <span class="n">count</span> <span class="o">+=</span> <span class="mi">1</span>
      <span class="k">except</span> <span class="n">Empty</span><span class="p">:</span>
        <span class="n">logging</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;dropped </span><span class="si">{0}</span><span class="s2"> items from queue&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">count</span><span class="p">))</span>
        <span class="n">done</span> <span class="o">=</span> <span class="kc">True</span></div></div>
</pre></div>

          </div>
        </div>
      </div>
      <div class="sphinxsidebar" role="navigation" aria-label="main navigation">
        <div class="sphinxsidebarwrapper">
<div id="searchbox" style="display: none" role="search">
  <h3>Quick search</h3>
    <div class="searchformwrapper">
    <form class="search" action="../../search.html" method="get">
      <input type="text" name="q" />
      <input type="submit" value="Go" />
      <input type="hidden" name="check_keywords" value="yes" />
      <input type="hidden" name="area" value="default" />
    </form>
    </div>
</div>
<script type="text/javascript">$('#searchbox').show(0);</script>
        </div>
      </div>
      <div class="clearer"></div>
    </div>
    <div class="related" role="navigation" aria-label="related navigation">
      <h3>Navigation</h3>
      <ul>
        <li class="right" style="margin-right: 10px">
          <a href="../../genindex.html" title="General Index"
             >index</a></li>
        <li class="right" >
          <a href="../../py-modindex.html" title="Python Module Index"
             >modules</a> |</li>
        <li class="nav-item nav-item-0"><a href="../../index.html">TensorFlowOnSpark 1.3.3 documentation</a> &#187;</li>
          <li class="nav-item nav-item-1"><a href="../index.html" >Module code</a> &#187;</li> 
      </ul>
    </div>
    <div class="footer" role="contentinfo">
        &#169; Copyright 2018, Yahoo Inc.
      Created using <a href="http://sphinx-doc.org/">Sphinx</a> 1.7.9.
    </div>
  </body>
</html>